#include <windows.h>
#include <winhttp.h>
#include <stdarg.h>
#include <intrin.h>
#include <stdbool.h>

typedef struct
{
    HANDLE hHeap;
    HANDLE hStdOut;
    HANDLE hStdErr;
} Handles;

void __forceinline Handles_init(Handles *hs)
{
    hs->hHeap = GetProcessHeap();
    hs->hStdOut = GetStdHandle(STD_OUTPUT_HANDLE);
    hs->hStdErr = GetStdHandle(STD_ERROR_HANDLE);
}

void Write(HANDLE h, LPVOID s, DWORD n) { WriteFile(h, s, n, &n, NULL); }
void PrintExA(HANDLE h, LPCSTR s, DWORD n) { WriteConsoleA(h, s, n, NULL, NULL); }
void PrintExW(HANDLE h, LPCWSTR s, DWORD n) { WriteConsoleW(h, s, n, NULL, NULL); }
void PrintfA(HANDLE h, LPCSTR fmt, ...)
{
    va_list ap;
    va_start(ap, fmt);
    LPSTR s;
    DWORD dwSize = FormatMessageA(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_STRING, fmt, 0, 0, (LPSTR)&s, 0, &ap);
    if (s)
    {
        PrintExA(h, s, dwSize);
        LocalFree(s);
    }
    va_end(ap);
}
void PrintfW(HANDLE h, LPCWSTR fmt, ...)
{
    va_list ap;
    va_start(ap, fmt);
    LPWSTR s;
    DWORD dwSize = FormatMessageW(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_STRING, fmt, 0, 0, (LPWSTR)&s, 0, &ap);
    if (s)
    {
        PrintExW(h, s, dwSize);
        LocalFree(s);
    }
    va_end(ap);
}

#define PrintA(h, s) PrintExA(h, s, sizeof(s) - 1)
#define PrintW(h, s) PrintExW(h, s, sizeof(s) - 1)
#define Print(h, s) PrintEx(h, s, sizeof(s) - 1)

#ifndef _DEBUG
#define ERR(fun, code) ((void)0)
#define E(fun) ((void)0)
#else
#define ERR(fun, code) ErrorMessage(TEXT(__FILE__), __LINE__, TEXT(__FUNCSIG__), TEXT(fun), code)
#define E(fun)                     \
    do                             \
    {                              \
        DWORD dw = GetLastError(); \
        if (dw)                    \
            ERR(fun, dw);          \
    } while (0)
void ErrorMessage(LPCTSTR lpszFileName, DWORD dwLine, LPCTSTR lpszParentFunction, LPCTSTR lpszFunction, DWORD dwCode); // 输出错误信息
#endif

#if defined(UNICODE) || defined(_UNICODE)
#define PrintEx PrintExW
#define Printf PrintfW
#else
#define PrintEx PrintExA
#define Printf PrintfA
#endif

void mainCRTStartup()
{
    Handles hs;
    Handles_init(&hs);
    int argc;
    LPCWSTR *argv = CommandLineToArgvW(GetCommandLineW(), &argc);
    if (!argv)
        E("CommandLineToArgvW");
    else
    {
        if (argc < 2 || argc > 3)
            PrintfW(hs.hStdErr, L"usage: %s url [output filepath]\n", argv[0]);
        else
        {
            HANDLE hFile = hs.hStdOut;
            if (argc == 3)
            {
                hFile = CreateFileW(argv[2], GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
                if (INVALID_HANDLE_VALUE == hFile)
                    E("CreateFileW");
            }
            if (hFile != INVALID_HANDLE_VALUE)
            {
                HINTERNET hOpen = WinHttpOpen(NULL, WINHTTP_ACCESS_TYPE_AUTOMATIC_PROXY, NULL, NULL, 0);
                if (!hOpen)
                    E("WinHttpOpen");
                else
                {
                    URL_COMPONENTS url_c;
                    __stosb((PBYTE)&url_c, 0, sizeof url_c);
                    url_c.dwStructSize = sizeof url_c;
                    url_c.dwHostNameLength = 1;
                    url_c.dwExtraInfoLength = 1;
                    url_c.dwUrlPathLength = 1;
                    if (!WinHttpCrackUrl(argv[1], 0, 0, &url_c))
                        E("WinHttpCrackUrl");
                    else
                    {
                        LPWSTR hostName = NULL;
                        DWORD dwFlag = 0;
                        HINTERNET hConnect;
                        DWORD dwErrCode;
                        if (url_c.lpszHostName)
                        {
                            hostName = HeapAlloc(hs.hHeap, 0, sizeof(WCHAR) * (url_c.dwHostNameLength + 1));
                            __movsb((PBYTE)hostName, (const BYTE *)url_c.lpszHostName, url_c.dwHostNameLength * sizeof(WCHAR));
                            hostName[url_c.dwHostNameLength] = 0;
                        }
                        switch (url_c.nScheme)
                        {
                        case INTERNET_SCHEME_HTTPS:
                            dwFlag = WINHTTP_FLAG_SECURE;
                            break;
                        }
                        hConnect = WinHttpConnect(hOpen, hostName, url_c.nPort, 0);
                        dwErrCode = GetLastError();
                        if (hostName)
                            HeapFree(hs.hHeap, 0, hostName);
                        if (!hConnect)
                            ERR("WinHttpConnect", dwErrCode);
                        else
                        {
                            HINTERNET hRequest;
                            LPWSTR objName = NULL;
                            DWORD objNameLen = url_c.dwUrlPathLength + url_c.dwExtraInfoLength;
                            if (objNameLen)
                            {
                                objName = HeapAlloc(hs.hHeap, 0, sizeof(WCHAR) * (objNameLen + 1));
                                __movsb((PBYTE)objName, (const BYTE *)url_c.lpszUrlPath, url_c.dwUrlPathLength * sizeof(WCHAR));
                                __movsb((PBYTE)(objName + url_c.dwUrlPathLength), (const BYTE *)url_c.lpszExtraInfo, url_c.dwExtraInfoLength * sizeof(WCHAR));
                                objName[objNameLen] = 0;
                            }
                            hRequest = WinHttpOpenRequest(hConnect, NULL, objName, NULL, WINHTTP_NO_REFERER, WINHTTP_DEFAULT_ACCEPT_TYPES, dwFlag);
                            dwErrCode = GetLastError();
                            if (objName)
                                HeapFree(hs.hHeap, 0, objName);
                            if (!hRequest)
                                ERR("WinHttpOpenRequest", dwErrCode);
                            else
                            {
                                if (!WinHttpSendRequest(hRequest, WINHTTP_NO_ADDITIONAL_HEADERS, 0, WINHTTP_NO_REQUEST_DATA, 0, 0, 0))
                                    E("WinHttpSendRequest");
                                else if (!WinHttpReceiveResponse(hRequest, NULL))
                                    E("WinHttpReceiveResponse");
                                else
                                {
                                    DWORD dwSize = 0;
                                    do
                                    {
                                        if (!WinHttpQueryDataAvailable(hRequest, &dwSize))
                                            E("WinHttpQueryDataAvailable");
                                        else
                                        {
                                            char *buffer = HeapAlloc(hs.hHeap, 0, dwSize);
                                            DWORD dwRead;
                                            if (!WinHttpReadData(hRequest, buffer, dwSize, &dwRead))
                                            {
                                                E("WinHttpReadData");
                                                dwSize = 0;
                                            }
                                            else
                                                Write(hFile, buffer, dwRead);
                                            HeapFree(hs.hHeap, 0, buffer);
                                        }
                                    } while (dwSize);
                                    if (argc == 2)
                                        PrintA(hs.hStdOut, "\n");
                                }
                                if (!WinHttpCloseHandle(hRequest))
                                    E("WinHttpCloseHandle");
                            }
                            if (!WinHttpCloseHandle(hConnect))
                                E("WinHttpCloseHandle");
                        }
                        if (!WinHttpCloseHandle(hOpen))
                            E("WinHttpCloseHandle");
                    }
                    if (argc == 3)
                        CloseHandle(hFile);
                }
            }
        }
        LocalFree(argv);
    }
    ExitProcess(0);
}

#ifdef _DEBUG
void ErrorMessage(LPCTSTR lpszFileName, DWORD dwLine, LPCTSTR lpszParentFunction, LPCTSTR lpszFunction, DWORD dwCode)
{
    LPTSTR lpMsgBuf = NULL;
    LPTSTR lpDisplayBuf = NULL;
    DWORD dwSize = 0;

    FormatMessage(
        FORMAT_MESSAGE_ALLOCATE_BUFFER |
            FORMAT_MESSAGE_FROM_SYSTEM |
            FORMAT_MESSAGE_IGNORE_INSERTS |
            FORMAT_MESSAGE_MAX_WIDTH_MASK,
        NULL,
        dwCode,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPTSTR)&lpMsgBuf,
        0, NULL);

#ifdef __cplusplus
    using va_list_array = va_list[];
#else
#define va_list_array (va_list[])
#endif

    if (lpMsgBuf)
        dwSize = FormatMessage(
            FORMAT_MESSAGE_ALLOCATE_BUFFER |
                FORMAT_MESSAGE_ARGUMENT_ARRAY |
                FORMAT_MESSAGE_FROM_STRING,
            TEXT("Filename:%1\n")
                TEXT("Line:%2!lu!\n")
                    TEXT("Parent function:%3\n")
                        TEXT("Function:%4\n")
                            TEXT("Code:%5!lu!\n")
                                TEXT("Description:%6\n"),
            0,
            0,
            (LPTSTR)&lpDisplayBuf,
            0,
            va_list_array{
                (va_list)lpszFileName,
                (va_list)(SIZE_T)dwLine,
                (va_list)lpszParentFunction,
                (va_list)lpszFunction,
                (va_list)(SIZE_T)dwCode,
                (va_list)lpMsgBuf});

    // Display the error message and exit the process
    if (lpDisplayBuf && dwSize)
        WriteConsole(GetStdHandle(STD_ERROR_HANDLE), lpDisplayBuf, dwSize, NULL, NULL);

    if (lpMsgBuf)
        LocalFree(lpMsgBuf);
    if (lpDisplayBuf)
        LocalFree(lpDisplayBuf);
}
#endif